import torch
import numpy as np
from pathlib import Path
from src.models.refiner import ROIRefinerModel
from src.config import DEVICE


class EnsembleInference:
    """
    

    
    1. Weighted Average
    2. Majority Voting
    3. NMSWeighted Boxes Fusion
    """

    def __init__(self, model_paths, weights=None, strategy='weighted_average'):
        """
        Args:
            model_paths: 
            weights: None=
            strategy: 'weighted_average', 'voting', 'wbf'
        """
        print(f"  ({len(model_paths)})...")

        self.models = []
        for i, path in enumerate(model_paths):
            model = ROIRefinerModel(device=DEVICE)
            checkpoint = torch.load(path, map_location=DEVICE)
            if 'model_state_dict' in checkpoint:
                model.load_state_dict(checkpoint['model_state_dict'])
            else:
                model.load_state_dict(checkpoint)
            model.eval()
            self.models.append(model)
            print(f"     {i + 1} ")

        self.weights = weights if weights else [1.0] * len(model_paths)
        self.weights = np.array(self.weights) / sum(self.weights)
        self.strategy = strategy

    @torch.no_grad()
    def predict(self, roi_batch):
        """
        

        Returns:
            class_logits, bbox_deltas
        """
        all_cls_logits = []
        all_bbox_deltas = []

        # 
        for model in self.models:
            cls, reg = model(roi_batch)
            all_cls_logits.append(cls)
            all_bbox_deltas.append(reg)

        # 
        if self.strategy == 'weighted_average':
            # 
            cls_logits = sum(w * logits for w, logits in zip(self.weights, all_cls_logits))
            bbox_deltas = sum(w * deltas for w, deltas in zip(self.weights, all_bbox_deltas))

        elif self.strategy == 'voting':
            # 
            cls_probs = [torch.softmax(logits, dim=1) for logits in all_cls_logits]
            cls_preds = [probs.argmax(dim=1) for probs in cls_probs]

            # 
            stacked = torch.stack(cls_preds, dim=0)  # [num_models, batch]
            cls_logits = torch.mode(stacked, dim=0).values

            # 
            bbox_deltas = sum(self.weights[i] * deltas for i, deltas in enumerate(all_bbox_deltas))

        elif self.strategy == 'wbf':
            # WBF
            cls_logits = sum(w * logits for w, logits in zip(self.weights, all_cls_logits))
            bbox_deltas = sum(w * deltas for w, deltas in zip(self.weights, all_bbox_deltas))

        return cls_logits, bbox_deltas


# 
if __name__ == '__main__':
    from src.inference.local_inference import LocalInference

    # 
    model_paths = [
        'weights/stage2_refiner.pth',
        'weights/stage2_refiner_ema.pth',
        'weights/stage2_refiner_swa.pth'
    ]

    # 
    ensemble = EnsembleInference(
        model_paths=[p for p in model_paths if Path(p).exists()],
        weights=[1.0, 1.2, 1.0],  # EMA
        strategy='weighted_average'
    )

    # 
    # inferencer.refiner = ensemble